import numpy as np
from PIL import ImageGrab
import cv2
import time

import Tools.grabscreen as grabscreen
import Tools.getkeys as getkeys
import os

lower_hsv_r = np.array([8, 215, 134])
upper_hsv_r = np.array([19, 253, 255])
lower_hsv_b = np.array([0, 0, 255])
upper_hsv_b = np.array([109, 23, 255])
wait_time = 3
L_t = 3
file_name = 'train_data.npy'  # 人工训练时保存操作的文件
window_size = (0, 100, 640, 448)  # 截取屏幕的大小，四个参数分别是x起点，y起点 x终点，y终点

if os.path.isfile(file_name):
    print("指定数据文件存在，加载完成！")
    training_data = list(np.load(file_name, allow_pickle=True))
else:
    print("指定数据文件不存在，已建立新文件！")
    training_data = []

for i in list(range(wait_time))[::-1]:
    print(i + 1)
    time.sleep(1)

last_time = time.time()
while (True):
    output_key = getkeys.get_key(getkeys.key_check())  # 按键收集
    if output_key == [1, 1, 1, 1, 1, 1]:
        print(len(training_data))
        np.save(file_name, training_data)
        break

    hsv = cv2.cvtColor(grabscreen.grab_screen(window_size), cv2.COLOR_BGR2HSV)  # HSV收集
    mask1 = cv2.inRange(hsv, lower_hsv_r, upper_hsv_r)
    mask2 = cv2.inRange(hsv, lower_hsv_b, upper_hsv_b)
    mask = cv2.bitwise_xor(mask1, mask2)
    screen_reshape = cv2.resize(mask, (160, 86))  # 压缩图像大小
    training_data.append([screen_reshape, output_key])

    if len(training_data) % 1000 == 0:  # 每500次loop保存一次以免异常情况出现
        print(len(training_data))
        np.save(file_name, training_data)

    cv2.imshow('cenTest', mask)  # 显示处理完的图像
    cv2.moveWindow('cenTest', 1000, 200)  # 移动到指定位置，因人调整

    # 测试帧数用
    print('抓帧耗时 {} 秒'.format(time.time() - last_time))
    last_time = time.time()

    if cv2.waitKey(5) & 0xFF == ord('q'):
        break
cv2.waitKey()  # 视频结束后，按任意键退出
cv2.destroyAllWindows()
